import random

import numpy as np
import matplotlib.pyplot as plt
import math
import matplotlib.animation as anime  # use this to generate GIF
import matplotlib.patches as patches # use this to draw geometry
from position import *

def dist(x1, y1, x2, y2):
    """

    :param x1: x1
    :param y1: y1
    :param x2: x2
    :param y2: y2
    :return: distance
    """
    return math.sqrt((x1 - x2) ** 2 + (y1 - y2) ** 2)


### Iteration One: Suppose obstacles and robots are circular. they are no longer mass dots with no morphology

# ! further requirement: make agent and obstacles rectangles, instead of circles.
# ! use a modern bump detection algorithm to efficiently implement it. warp this algorithm into a func or API

# !! even more, may let the size and morphology of robots change over time, as the robot may strech its gripper,
# !! when it is ready to kick the ball, which leads to a noticable change to its bounding box

#*** global variables ********************************

# add one attribute: radius

# robots are considered as homogeneous now, so r_agent and r_obs should be set as the same
# ! further requirement: adapt to heterogeneous case
r_agent = 0.15
r_obs = 0.15

# there is no need to modify the potential energy expr
# however, we need to re-define repulsive energy field, or pass a reaonable 'radius' parameter to the original one
# let us set a global warning coefficent that multiplies with object radius to define the range of repulsive field
k_warn = 2

# then, calc the distance when repulsive force takes effect
r_rep = k_warn * r_obs

# let us define another coefficient that is used to cap the repulsive energy when dist_gap is very close to 0
# this prevents 1/dist_gap to diverge when dist_gap approaches zero, which leads to unwanted math errors "ZeroDivisionError"
k_cap = 1.1

# use k_cap to set the maximum 1/r value. denote this r as r_cap
r_cap = k_cap * r_obs

## notice this: r_obs < r_cap < r_rep ; r_agent
## r_obs < r_cap = r_obs * 1.2 < r_rep = r_obs * 1.5

# set a range that determines whether agent is close enough to target, so that we need to keep the speed not decreasing too much
# notice that this outlines a rectangle, not a circle, as it limits the lower bound of the deterrance value along both x and y axes
r_tar = 0.25

# ! to implement: set a variable to maintain whether agent collides with any obstcale currently
# if no bumped obstacles at this timestep, it is set to blank.
# if bumped, it contains tuples that records the details of collision:
# every collision is expressed in a tuple formatted as follows:
# (index of obstacle -> int; position x of agent, y of agent, position x of obstacle, y of obstacle; distance between their centers -> float)
# example: (3  ,1.23,4.66,1.10,4.50,  0.35)
if_bump = False
bump_detail_latest = []

# *** end of global variables ***************************

# define a new dist function to determine the minimum distance between two objects' surfaces
def dist_gap(x1,y1,x2,y2,ra,ro):
    """

    :param x1: x of agent
    :param y1: y of agent
    :param x2: x of obstacle
    :param y2: y of obstacle
    :param ra: radius of agent
    :param ro: radius of obstacle
    :return: minimum gap between them. if bumped together, return 0
    """
    return max(dist(x1,y1,x2,y2)-ra-ro,0)

# modify the defualt parameter of potential energy: modify rep_radius = r_rep
def potential_energy(current_x, current_y, endx, endy, obs_x_collection, obs_y_collection, coeff_att=1, coeff_rep=10,
                     rep_radius=r_rep):
    """

    :param current_x: start point x
    :param current_y: start point y
    :param endx: end point x
    :param endy: end point y
    :param obs_x_collection: a list containing the obstacle's x
    :param obs_y_collection: a list containing the obstacle's y, with responds to x_collection
    :param coeff_att: coefficient of attractive energy
    :param coeff_rep: coefficient of repulsive energy
    :param rep_radius: maximum repulsion range
    :return:
    """

    def attract_energy(xx, yy, ex, ey, coeff=coeff_att):
        """

        :param xx:  current x
        :param yy:  current y
        :param ex:  target x
        :param ey:  target y
        :param coeff: coefficient for energy
        :return: potential energy of attractive energy
        """
        r = dist(xx, yy, ex, ey)
        return coeff * r ** 2

    # alse modify default value of rep_range as r_rep
    def repulse_energy(xx, yy, ox_list, oy_list, rep_range=r_rep, coeff=coeff_rep):
        """
        take care! this function directly uses global variables:

        r_cap, r_agent, r_obs.
        without referring them as function default input params
        it takes r_rep , another global var, as its input parameter, instead

        :param xx: current x
        :param yy: current y
        :param ox_list: obstacle's x list
        :param oy_list: obstacle's y list
        :param rep_range: maximum repulsion range, set to r_rep by default
        :param coeff: coefficient for repulsive energy
        :return:
        """

        ans = 0
        for ox, oy in zip(ox_list, oy_list):
            # r = dist(xx, yy, ox, oy)
            # change "r" here to the representation of minimum surface distance
            r = dist_gap(xx, yy, ox, oy, r_agent,r_obs)

            energy = 0
            if r <= rep_range:
                if r > r_cap:  # in repulsion range but haven't capped
                    energy = coeff * (1 / r - 1 / rep_range) ** 2
                elif r > 0 and r <= r_cap: # in repulsion range, capped, but haven't bumped
                    energy = coeff * (1 / r_cap - 1 / rep_range) ** 2
                else: # have already bumped together !!!
                # ! to implement: save bump details to global var 'bump_detail_latest'
                # ! in a form of tuple defined above.   list[tuple(),...]
                # ! if there are multiple bumps at one timestep, append all tuples into the list
                # ! remember to clear restore 'bump_detail_latest' to empty and 'if_bumped' to False when a new iter starts
                    energy = 10000 # set the energy when bumped together. should be big enough
            ans += energy
        return ans

    e = 0

    e_att = attract_energy(current_x, current_y, endx, endy, coeff_att)
    e += e_att

    e_rep = repulse_energy(current_x, current_y, obs_x_collection, obs_y_collection, rep_radius, coeff_rep)
    e += e_rep

    return e

# also modify the gradient expr:
# set default param of rep_radius as r_rep
# since we have introduced 'warning range' r_rep and 'energy cap range' r_cap,
# we also need to modify the gradient expr into a form of segmented function
def gradient(current_x, current_y, endx, endy, obs_x_collection, obs_y_collection, coeff_att=1, coeff_rep=10,
             rep_radius=r_rep):
    """
    take care! this function directly uses global variables mentioned above
    without referring them as function default input params:
    r_cap, r_agent, r_obs, r_tar
    notice that it takes r_tar into account,additionally

    it takes r_rep , another global var, as its input parameter named 'rep_radius',

    :param current_x: start point x
    :param current_y: start point y
    :param endx: end point x
    :param endy: end point y
    :param obs_x_collection: a list containing the obstacle's x
    :param obs_y_collection: a list containing the obstacle's y, with responds to x_collection
    :param coeff_att: coefficient of attractive energy
    :param coeff_rep: coefficient of repulsive energy
    :param rep_radius: maximum repulsion range
    :return:a tuple of gradient along x and y-axis
    """

    # actually, there is no definite need to modify attract gradient

    # however, if we takes gradient as a indicator of speed, we won't want speed to decrease to a minor value when close to target.
    # we need to guarantee a minimum speed in order to keep the robot running at a feasible speed towards target, no matter where it is.
    # so we also express it as a segmented function --  we add a threshold to minimum speed.

    # in another sense, we also hope the speed not too large -- we add another threshold to maximum speed.(! haven't implemented yet)

    # in order to make these two boundaries effectively regulate speed, not making it always at minimum or maximum value
    # wherever the agent is. we need to understand the possible value range of gradient.

    def gradient_attract(xx, yy, ex, ey, coeff=coeff_att):
        """
        this function uses global var: r_tar
        :param xx:  current x
        :param yy:  current y
        :param ex:  target x
        :param ey:  target y
        :param coeff: coefficient for energy
        :return: a tuple of gradient along x and y axis
        """
        # we use the original r instead of r_dist, which substracts r_obs and r_agent from r.
        # actually, as the target is only a point, without geometric covering area, we shouldn't use r_obs
        # we need to keep the gradient at a constant minimum value when agent is so 'close' to the target that
        # the original formula gives a value very close to zero, leading to an unreachable final position

        r = dist(xx, yy, ex, ey)


        # add a minimum speed and maximum speed constraint.
        # ! to implement: maximum value haven't set yet. set minimum value first
        # now, directly use the global variable 'r_tar'

        # remember r_tar is the absolute value. should discuss about sign and negativity
        xe = xx - ex
        ye = yy - ey
        if xe >=0: #位矢为正
            rx = max(r_tar,xe) / r  # 单位矢量在x方向的投影, r_tar
        else: #位矢为负
            rx = min(-r_tar,xe) / r

        if ye >=0:
            ry = max(r_tar,ye) / r  # 单位矢量在y方向的投影
        else:
            ry = min(-r_tar,ye) / r

        gradient_x = 2 * coeff * r * rx  # x方向梯度
        gradient_y = 2 * coeff * r * ry  # y方向梯度

        return gradient_x, gradient_y

    # modify repulse grads according to morphology of obstacles:
    # set the default value of 'radius' into global var, r_rep
    def gradient_repulse(xx, yy, ox_list, oy_list, radius=r_rep, coeff=coeff_rep):
        """
        this function uses global var: r_obs, r_agent, r_cap.
        r_rep is passed through a variable
        :param xx: current x
        :param yy: current y
        :param ox_list: obstacle's x list
        :param oy_list: obstacle's y list
        :param radius: maximum repulsion range
        :param coeff: coefficient for repulsive energy
        :return: a tuple of gradient along x and y axis
        """
        ans_x = 0
        ans_y = 0
        grd_x = 0
        grd_y = 0
        for ox, oy in zip(ox_list, oy_list):
            r = dist(xx, yy, ox, oy)

            #calc gradient contribution of every obstacle
            if r <= radius: # consider repulsion. r_cap is used to limit grads above lowest threshold

                # remember r_tar is the absolute value. should discuss about sign and negativity
                xox = xx - ox
                yoy = yy - oy

                if xox >=0:
                    rx = max(r_cap, xox) / r  # 单位矢量在x方向的投影, r_cap
                else:
                    rx = min(-r_cap, xox)

                if yoy >=0:
                    ry = max(r_cap, yoy) / r  # 单位矢量在y方向的投影
                else:
                    ry = min(-r_cap, yoy)

                grd_x = 2 * coeff * (1 / r - 1 / radius) ** (-r ** 2) * rx
                grd_y = 2 * coeff * (1 / r - 1 / radius) ** (-r ** 2) * ry
            else: # r > radius, no repulsive force
                grd_x,grd_y=0,0

            ans_x += grd_x
            ans_y += grd_y
        return ans_x, ans_y

    g_att_x, g_att_y = gradient_attract(current_x, current_y, endx, endy, coeff_att)
    g_rep_x, g_rep_y = gradient_repulse(current_x, current_y, obs_x_collection, obs_y_collection, rep_radius, coeff_rep)
    grdx = g_att_x + g_rep_x
    grdy = g_att_y + g_rep_y
    return grdx, grdy


# this function regulate the 'learning rate', in other words, another multiplier to determine speed according to grads
def step_size(current_x, current_y, start_x, start_y, end_x, end_y, max_step_size=0.1, min_step_size=0.05, decay=1):

    r = dist(current_x, current_y, end_x, end_y)
    r_ref = dist(start_x, start_y, end_x, end_y)

    ratio = min(r / r_ref,1)
    step = max_step_size * math.exp(-decay * (1 - ratio))

    if step < min_step_size:
        return min_step_size
    return step

# one iteration of decision making
def step(current_x,current_y,start_x,start_y,end_x,end_y,obs_x_collection,obs_y_collection,
         coeff_att=1,coeff_rep=10,rep_radius=0.1,
         max_step_size=0.1,min_step_size=0.05,decay=1):

    stride = step_size(current_x,current_y,start_x,start_y,end_x,end_y,max_step_size,min_step_size,decay) # current "lr"
    grd_x,grd_y=gradient(current_x,current_y,end_x,end_y,obs_x_collection,obs_y_collection,coeff_att,coeff_rep,rep_radius)
    energy = potential_energy(current_x,current_y,end_x,end_y,obs_x_collection,obs_y_collection,coeff_att,coeff_rep,rep_radius)

    #update current_x and current_y
    current_x,current_y = current_x-stride*grd_x,current_y-stride*grd_y

    return current_x,current_y,energy,grd_x,grd_y,stride   # return new position and energy and gradient

### Visualization
# generate a frame describing agent, start point, target point and all obstacles
def draw(startx,starty,endx,endy,positions,obs_x_list, obs_y_list,suceeded =False,max_points = 20, set_range = True):
    fig, ax = plt.subplots()
    plt.xlabel('length')
    plt.ylabel('width')
    plt.title('simulation')

    # determine graph range
    # get range of obstacles

    if set_range:
        margin = 0.8
        obs_xmax, obs_xmin = max(obs_x_list) + margin, min(obs_x_list) - margin
        obs_ymax, obs_ymin = max(obs_y_list) + margin, min(obs_y_list) - margin

        # compare with start and end

        xmin = min(obs_xmin, startx, endx)
        xmax = max(obs_xmax, startx, endx)
        ymin = min(obs_ymin, starty, endy)
        ymax = max(obs_ymax, starty, endy)

        plt.xlim(xmin, xmax)
        plt.ylim(ymin, ymax)
    else:
        plt.xlim(0, 10)
        plt.ylim(0, 8)

    metadata = dict(title="Movie", artist="sourabh")

    writer = anime.PillowWriter(fps=20, metadata=metadata)

    length = len(positions)
    num_obs = len(obs_x_list)

    xs = []
    ys = []
    index = 0

    if suceeded:
        allowed_index = np.arange(max_points).tolist()
        start = max_points
        max_points *=4
        for i in range(max_points):
            #cur_index = min(length-1 ,start + i/max_points*(length-start))
            cur_index = start + (length - start) * math.log(1 + i / max_points)
            cur_index = int(cur_index)
            if cur_index >= length:
                cur_index = length - 1
            allowed_index.append(cur_index)
        '''
        allowed_index = np.arange(0,length)
        allowed_index_ref = 1 - np.arange(1,max_points+1)/max_points
        allowed_index_exp = np.exp(allowed_index_ref)
        '''
        with writer.saving(fig,"result.gif",100):
            # draw start and end points as patches.Rectangle() objects
            # start point
            start_rect = patches.Rectangle((startx, starty), 0.5, 0.5, facecolor='r', edgecolor='rosybrown')
            ax.add_patch(start_rect)

            # end point
            end_rect = patches.Rectangle((endx, endy), 0.5, 0.5, facecolor='b', edgecolor='brown')

            # add these two points into patch using ax.add_patch()
            ax.add_patch(end_rect)

            # draw obstacles
            for i in range(0, num_obs):
                # circle to be drawn

                # center
                obs_x = obs_x_list[i]
                obs_y = obs_y_list[i]

                # radius
                rad = r_obs  # use global data: r_obs

                # draw an obstacle as patches.Circle() object
                obs_circle = patches.Circle((obs_x, obs_y), rad, facecolor='w', edgecolor='y')

                # add obstacle into current patch
                ax.add_patch(obs_circle)

            #for posx,posy in iter(positions): # iterate through all position points
            for j in allowed_index:
                posx,posy = positions[j]
                index+=1
                # get current agent position into posx and posy

                # draw current agent
                # draw agent as patches.Circle() object
                r_a = r_agent  # circle radius
                agent_circle = patches.Circle((posx, posy), r_a,facecolor='g',edgecolor='r') # use global var: r_agent

                # add obstacle into current patch
                ax.add_patch(agent_circle)


                # draw obstacles
                # for one timestep, draw every obstacle as a circle

                # flush current figure into gif file as a frame added to the end
                writer.grab_frame()

        # show pic generated
        plt.show()
    else:
        print("No picture given since path finding failed")
        # get current agent position into posx and posy

        # draw start and end points as patches.Rectangle() objects
        # start point
        start_rect = patches.Rectangle((startx, starty), 0.3, 0.3, facecolor='r', edgecolor='rosybrown')

        ax.add_patch(start_rect)
        # end point
        end_rect = patches.Rectangle((endx, endy), 0.5, 0.5, facecolor='g', edgecolor='rosybrown')

        # add these two points into patch using ax.add_patch()
        ax.add_patch(end_rect)

        # failed, not drawing current agent

        # draw obstacles
        # for one timestep, draw every obstacle as a circle
        for i in range(0, num_obs):
            # circle to be drawn

            # center
            obs_x = obs_x_list[i]
            obs_y = obs_y_list[i]

            # radius
            rad = r_obs  # use global data: r_obs

            # draw an obstacle as patches.Circle() object
            obs_circle = patches.Circle((obs_x, obs_y), rad, facecolor='w', edgecolor='b')

            # add obstacle into current patch
            ax.add_patch(obs_circle)
        #show end pos
        end_failed = patches.Rectangle(positions[-1], 0.3, 0.3, facecolor='pink', edgecolor='rosybrown')
        ax.add_patch(end_failed)
        # only shows obstruct distribution
        plt.show()

# integrate all figures into a GIF file

def run(start, end, obs):
    # !!! remember also change these global variables defined above, if necessary
    # global vars include: r_rep; r_cap; r_tar; r_obs; r_agent

    # define start and end point position
    startx = start.x
    starty = start.y
    endx = end.x
    endy = end.y

    # define obstacle positions
    # prebuilt cases
    
    obs_x = []
    obs_y = []

    for o in obs:
        obs_x.append(o.x)
        obs_y.append(o.y)
        
    print(obs_x)
    print(obs_y)

    # switch on to test a random setting
    randomize = False

    if randomize:

        #anchor points of start and end
        startx_anchor,starty_anchor =0,0
        endx_anchor,endy_anchor = 6,6

        # random offset for both
        max_offset_x,max_offset_y = 1,1

        # generate a random start and end point

        startx = startx_anchor+random.uniform(-max_offset_x,max_offset_x)
        starty = starty_anchor+random.uniform(-max_offset_y,max_offset_y)
        endx = endx_anchor + random.uniform(-max_offset_x,max_offset_x)
        endy = endy_anchor + random.uniform(-max_offset_y,max_offset_y)

        # generate the number of obstacle
        obs_num =random.randint(4,10)
        obs_x=[]
        obs_y=[]
        # determine the range of potential obstacles
        obs_x_min,obs_y_min = 0,0  # minimum x and y coordinates
        obs_range_x, obs_range_y = 6,6 # distribution range along x and y axis
        # obs_x_max,obs_y_max = obs_x_min+obs_range_x,obs_range_y+obs_y_min # calculated maximum x and y coordinates

        # determine the minimum allowed distance between any obstacle and start/end point
        protect_start = 1.0
        protect_end = 1.0

        # generate every obstacle
        generated = 0
        while generated < obs_num:
            randseed_x = random.random()
            randseed_y = random.random()

            # try generate a obstacle
            obs_xx = obs_x_min + randseed_x*obs_range_x
            obs_yy = obs_y_min + randseed_y*obs_range_y

            # check if it is too close to start and end
            # if so, re-generate until far enough
            dist_to_start = dist(obs_xx,obs_yy,startx,starty)
            dist_to_end = dist(obs_xx,obs_yy,endx,endy)

            if dist_to_start>=protect_start and dist_to_end>=protect_end:
                generated +=1
                obs_x.append(obs_xx)
                obs_y.append(obs_yy)

    # initialize -----
    current_x,current_y=startx,starty
    iters =0
    # storage containers
    strides = []
    gradients = []
    energies = []
    positions = []
    distances = []
    # ! to implement: add a checklist, recording whether bumped
    # bumps = []

    # specify max iterations
    max_iter = 1000
    # specify coefficients and propulsion range
    att_coeff = 12
    rep_coeff = 9
    rep_range = r_rep
    # specify max and min step range, and decay factor
    step_max = 0.005
    step_min = 0.001
    decay = 0.9

    # specify allowed missing range to end point
    episilon = 0.03

    # do one step, until Finished or Failed
    finished = False
    failed = False
    # report every * iterations:
    report_every = 10
    
    # step one pace
    curx,cury,ene,grx,gry,std=step(current_x,current_y,startx,starty,endx,endy,obs_x,obs_y,att_coeff,rep_coeff,rep_range,
                                step_max,step_min,decay)
    # return curx-current_x, cury-current_y
    
    while(not finished) and (not failed):
        # step one pace
        curx,cury,ene,grx,gry,std=step(current_x,current_y,startx,starty,endx,endy,obs_x,obs_y,att_coeff,rep_coeff,rep_range,
                                   step_max,step_min,decay)

        # update position and store result
        strides.append(std)
        current_x,current_y=curx,cury
        energies.append(ene)
        positions.append((current_x,current_y))
        gradients.append((grx,gry))

        # ! to implement: save bump details into 'bumps' if necessary.
        # ! remember to clear restore 'bump_detail_latest' to empty and 'if_bumped' to False when a new iter starts !!!

        # update iteration counter
        iters +=1

        # calc and store distance to end point
        dist_from_end = dist(current_x,current_y,endx,endy)
        distances.append(dist_from_end)
        # report checkpoints
        debug = False
        if iters%report_every == 0 and not debug:
            print("Iteration %d --"%iters,"Pos: %.3f,%.3f"%(current_x,current_y) )
        if debug:
            print("Iteration %d --"%iters,"Pos: %.3f,%.3f;"%(current_x,current_y), "gradient: %.3f, %.3f;"%(grx,gry),
                  "stride:%.3f;"%std,"distance: %.3f"%dist_from_end, "energy: %.3f;"%ene)
        # check if finished or failed
        if iters >max_iter:
            failed = True
        if dist_from_end <= episilon:
            finished = True

    #### to implement: visualization
    draw(startx,starty,endx,endy,positions,obs_x,obs_y,(not failed) and finished,7)


    # post-processing
    if failed:
        print("Failed after %d iterations!\n"%max_iter)
    else: # finished
        print("suceeded after %d iterations!\n"%iters)
    # plt.plot(distances)
